import torch
import pathlib
from collections import OrderedDict

NAME_MAPPING = {
    "share": "spatial_extractor",
    "features2fixed_1": "features2branch1",
    "features2fixed_2": "features2branch2",
    "features2fixed_3": "features2branch3",
    "features2fixed_4": "features2branch4",
    "classifier_1": "classifier1",
    "classifier_2": "classifier2",
    "classifier_3": "classifier3",
    "classifier_4": "classifier4",
}

ori_state_dict_path = pathlib.Path("/home/georgehu/Downloads/1.pth")

ori_state_dict = torch.load(ori_state_dict_path, map_location="cpu")
new_state_dict = OrderedDict()
for name, value in ori_state_dict.items():
    names = name.split(".")
    if names[0] in NAME_MAPPING:
        names[0] = NAME_MAPPING[names[0]]
    new_name = ".".join(names)
    new_state_dict[new_name] = value

torch.save(new_state_dict, "/home/georgehu/Documents/Projects/george-reimplement/medical/MTFiST/model/origin_pretrain.pth")

